package 剑指offer;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.TreeSet;

public class Offer40_JZ29 {
    public static int partition(int[] array, int start, int end) {
        int pivot = array[start];
        while (start < end) {
            while (start<end && array[end]>=pivot) {
                end--;
            }
            array[start] = array[end];
            while (start<end && array[start]<=pivot) {
                start++;
            }
            array[end] = array[start];
        }
        array[start] = pivot;
        return start;
    }

    // 快排法
    public static ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        int len = input.length;
        if (k>len || len==0 || k<=0) {
             return new ArrayList<>();
         }
        int start = 0;
        int end = len - 1;
        int index = partition(input, start, end);
        while (index != k-1) {
            if (index > k-1) {
                end = index-1;
                index = partition(input, start, end);
            } else {
                start = index+1;
                index = partition(input, start, end);
            }
        }
        ArrayList<Integer> list = new ArrayList<>();
        for (int i=0; i<k; i++) {
            list.add(input[i]);
        }
        return list;
    }

    // 大根堆法
    public static ArrayList<Integer> GetLeastNumbers_Solution2(int [] input, int k) {
        ArrayList<Integer> list = new ArrayList<>();
        if (input.length == 0 || k <= 0 || k > input.length) {
            return list;
        }
        PriorityQueue<Integer> queue = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2.compareTo(o1);
            }
        });
        for (int i=0; i< input.length; i++) {
            if (queue.size() == k) {
                if (queue.peek() > input[i]) {
                    queue.poll();
                    queue.add(input[i]);
                } else {
                    continue;
                }
            } else {
                queue.add(input[i]);
            }
        }
        while (!queue.isEmpty()) {
            list.add(queue.poll());
        }
        return list;
    }


    public static void main(String[] args) {
        int[] arr = {4,5,1,6,2,7,3,8};
        System.out.println(GetLeastNumbers_Solution2(arr, 4).toString());
    }
}
